#lang racket

(require "util/utils.scm")
(require "util/gensym.scm")
(provide pass-hoist)

;; input:
;;  tagged and closure converted source code
;;  (i.e. the output of pass-scope)
;;
;; output:
;;  an assoc list with two keys: definitions and code
;;  definitions is an assoc list mapping names to code labels
;;  code is an assoc list mapping labels to the captured environment, bindings/parameters, and code body

(define (pass-hoist defs)
  (let* ((co (make-collector))
         (do (make-collector))
         (i (box 0)))
    (for-each (lambda (def) (pass-hoist-def def co do i)) defs)
    `((definitions . ,(collector-get do))
      (code . ,(collector-get co)))))

(define (pass-hoist-def def co do i)
  (match def
    (`(define ,name ,body)
     (collect-object! co `((label ,name) () () ,(pass-hoist-exp body co do name (box 0))))
     (collect-object! do `(define ,name (label ,name))))
    (`(dont-define ,body)
     (let ((name `(top-exp ,(unbox i))))
       (set-box! i (+ 1 (unbox i)))
       (collect-object! co `((label . ,name) () () ,(pass-hoist-exp body co do name (box 0))))
       (collect-object! do `(execute (label . ,name)))))
    (else (error 'pass-hoist-def "malformed toplevel definition" def))))

(define (desugar-quote exp)
  (cond ((string? exp) `(string ,exp))
        ((atomic? exp) `(atomic ,exp))
        ((null? exp) `(atomic ,exp))
        ((symbol? exp) `(atomic ,exp))
        ((pair? exp)
         `(builtin cons
                   ,(desugar-quote (car exp))
                   ,(desugar-quote (cdr exp))))
        (else (error 'desugar-quote "?" exp))))

(define (pass-hoist-exp exp co do label i)
  (match exp
    (`(quote ,x) 
     (let* ((sym (dgensym 'const))
            (def `(define ,sym ,(desugar-quote x))))
       (pass-hoist-def def co do i)
       (pass-hoist-exp `(var glo ,sym) co do label i)))
    (`(string ,s)
     (let* ((sym (dgensym 'const))
            (def `(define ,sym (atomic ,s))))
       (pass-hoist-def def co do i)
       (pass-hoist-exp `(var glo ,sym) co do label i)))
    (`(atomic ,x) `(atomic ,x))
    (`(var ,type ,x) `(var ,type ,x))
    (`(app . ,exps) `(app . ,(mapply pass-hoist-exp exps co do label i)))
    (`(known-app ,i . ,exps) `(known-app ,i . ,(mapply pass-hoist-exp exps co do label i)))
    (`(begin . ,exps) `(begin . ,(mapply pass-hoist-exp exps co do label i)))
    (`(closure ,env ,resolv-env ,binds ,body)
     (let* ((counter (unbox i))
            (clo-label `(label ,(dgensym 'closure))))
       ;;`(label clo ,label ,counter)))
       (set-box! i (+ counter 1))
       (collect-object! co `(,clo-label ,env ,binds ,(pass-hoist-exp body co do label i)))
       `(closure ,resolv-env ,clo-label)))
    (`(if ,t ,c ,a)
     `(if ,(pass-hoist-exp t co do label i)
          ,(pass-hoist-exp c co do label i)
          ,(pass-hoist-exp a co do label i)))
    (`(let* ,bindings ,body)
     `(let* ,(map* (lambda (bind)
                     (list (car bind) (pass-hoist-exp (cadr bind) co do label i))) bindings)
            ,(pass-hoist-exp body co do label i)))
    (`(builtin ,nm . ,args)
     `(builtin ,nm . ,(mapply pass-hoist-exp args co do label i)))
    (else (error 'pass-hoist-exp "?" exp))))
